import os
import sqlite3

from opt.experiment import Experiment as BaseExperiment

from stacking.instance import Instance

            
class Experiment(BaseExperiment):
    def __init__(self, methods, directory, results="results.db"):
        BaseExperiment.__init__(self)
        self.methods = list(methods)
        self.directory = os.path.abspath(directory)
        if not os.path.isabs(results):
            results = os.path.join(directory, results)
        self.db = ResultsDB(filepath=results)
        
    def on_setup(self):
        self.db.connect()
        self.db.cursor.execute("CREATE TABLE IF NOT EXISTS results "
                               "(id REFERENCES completed(rowid),"
                               " objective REAL, cpu REAL, PRIMARY KEY (id))")
        self.db.mods_since_commit = 0
        
    def on_finalize(self):
        self.db.close()
        
    def iter_instances(self):
        for dirpath, dirnames, filenames in os.walk(self.directory):
            for filename in filenames:
                if filename.endswith(".sp"):
                    filepath = os.path.join(dirpath, filename)
                    if self.db.get_completed(filepath):
                        self.log.warning("SKIPPED: %s" % filepath)
                        continue
                    yield Instance.load(filepath)
                    self.db.set_completed(filepath)
                
    def iter_methods(self):
        instance_filepath = self.instance.meta.filepath
        for method in self.methods:
            if self.db.get_completed(instance_filepath, method.__name__):
                self.log.warning("SKIPPED: %s - %s" % (instance_filepath, method.__name__))
                continue
            yield method
            self.db.set_completed(instance_filepath, method.__name__)
            
    def iter_replications(self):
        instance_filepath = self.instance.meta.filepath
        method_name = self.method.__name__
        for i in xrange(self.replications):
            if self.db.get_completed(instance_filepath, method_name, i):
                self.log.warning("SKIPPED: %s - %s - %s" % (instance_filepath, method_name, i))
                continue
            yield i
            
    def save_results(self, solutions):
        instance_filepath = self.instance.meta.filepath
        method_name = self.method.__name__
        replication = self.replication
        self.log("Saving results for %s - %s - %s" % (instance_filepath, method_name, replication))
        self.db.set_completed(instance_filepath, method_name, replication)
        for solution in solutions:
            self.db.cursor.execute("INSERT INTO results " 
                                   "SELECT rowid, ?, ? FROM completed WHERE" 
                                   " instance=? AND method=? AND replication=?", 
                                   (solution.meta.objective, solution.meta.cpu, 
                                    instance_filepath, method_name, replication))
            if self.db.cursor.rowcount < 1:
                raise Exception("no data saved (no valid entry found in 'completed' table?)")
        self.db.mods_since_commit += 1
        if self.db.mods_since_commit >= 100:
            self.log.info("Committing results into database...")
            self.db.connection.commit()
            self.db.mods_since_commit = 0
            
            
class ResultsDB(object):
    def __init__(self, filepath=None):
        self.filepath = filepath
        self.connection = None
        self.cursor = None
        
    def connect(self, filepath=None):
        if filepath is not None:
            self.filepath = filepath
        if self.filepath is None:
            raise Exception("no file path provided for database")
        self.connection = sqlite3.connect(self.filepath)
        self.cursor = self.connection.cursor()
        self.init()
        
    def init(self):
        self.cursor.execute("CREATE TABLE IF NOT EXISTS completed "
                            "(instance TEXT, method TEXT, replication TEXT,"
                            " UNIQUE (instance, method, replication))")
        
    def close(self):
        self.cursor.close()
        self.connection.commit()
        self.connection.close()
        
    def get_completed(self, instance=None, method=None, replication=None):
        self.cursor.execute("SELECT 1 FROM completed WHERE"
                            " instance=? AND method=? AND replication=?", 
                            self._parse_args_completed(instance, method, replication))
        return self.cursor.fetchone() is not None
        
    def set_completed(self, instance=None, method=None, replication=None):
        self.cursor.execute("INSERT INTO completed VALUES(?, ?, ?)", 
                            self._parse_args_completed(instance, method, replication))
        
    def _parse_args_completed(self, instance, method, replication):
        if instance is None:
            assert method is None
            instance = "*"
        if method is None: 
            assert replication is None
            method = "*"
        if replication is None:
            replication = "*"
        return str(instance), str(method), str(replication)    
        
        
# --------------------------------------------------------------------------------------------------
# -- Methods
# --------------------------------------------------------------------------------------------------
from stacking.solvers import HeuristicSolver
from stacking.solvers.heuristics import RAND, CM, FO, PFO

def make_heur_solver(heuristic):
    def method(**params):
        solver = HeuristicSolver(heuristic)
        solver.verbose = False
        solver.solve(params["instance"], seed=params["replication"])
        return [solver.solution]
    method.__name__ = heuristic.__name__
    return method
    
methods = list()
for heuristic in [RAND, CM, FO, PFO]:
    methods.append(make_heur_solver(heuristic))
    
    